{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Polynomial Fit Plot\n",
    "===================\n",
    "This example shows how to overlay data with a fitted polynomial\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import altair as alt\n",
    "alt.data_transformers.enable('json')\n",
    "\n",
    "# Generate some random data\n",
    "rng = np.random.RandomState(1)\n",
    "x = rng.rand(40) ** 2\n",
    "y = 10 - 1. / (x + 0.1) + rng.randn(40)\n",
    "df = pd.DataFrame({'x': x, 'y': y})\n",
    "\n",
    "# Define the degree of the polynomial fit\n",
    "degree_list = [1, 3, 5]\n",
    "\n",
    "# Build a dataframe with the fitted data\n",
    "poly_data = pd.DataFrame({'xfit': np.linspace(df['x'].min(), df['x'].max(), 500)})\n",
    "\n",
    "for degree in degree_list:\n",
    "    poly_data[str(degree)] = np.poly1d(np.polyfit(df['x'], df['y'], degree))(poly_data['xfit'])\n",
    "\n",
    "# Tidy the dataframe so 'degree' is a variable\n",
    "poly_data = pd.melt(poly_data,\n",
    "                    id_vars=['xfit'],\n",
    "                    value_vars=[str(deg) for deg in degree_list],\n",
    "                    var_name='degree', value_name='yfit')\n",
    "\n",
    "# Plot the data points on an interactive axis\n",
    "points = alt.Chart(df).mark_circle(color='black').encode(\n",
    "    x=alt.X('x', axis=alt.Axis(title='x')),\n",
    "    y=alt.Y('y', axis=alt.Axis(title='y')),\n",
    ").interactive()\n",
    "\n",
    "# Plot the best fit polynomials\n",
    "polynomial_fit = alt.Chart(poly_data).mark_line().encode(\n",
    "    x='xfit',\n",
    "    y='yfit',\n",
    "    color='degree'\n",
    ")\n",
    "\n",
    "points + polynomial_fit\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "env": {},
   "interrupt_mode": "signal",
   "language": "python",
   "metadata": {},
   "name": "python3"
  },
  "language": "python"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
